import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt


# 假设你有一个简单的神经网络模型定义如下
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)  # 确保这里的输入维度和输出维度与保存的模型一致
        self.fc2 = nn.Linear(128, 10)  # 确保这里的输入维度和输出维度与保存的模型一致

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


# 定义数据转换
transform = transforms.Compose([
    transforms.Grayscale(),  # 确保图像是灰度图
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 初始化模型
model = SimpleNN()

# 加载模型状态字典
# model.load_state_dict(torch.load('./build/mnist_simplenn.pth', map_location=torch.device('cpu')))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f'Using device: {device}')
model.load_state_dict(torch.load('./build/mnist_simplenn.pth', map_location=device, weights_only=True))

# 设置模型为评估模式
model.eval()

# 加载并预处理单张图片
# image_path = './model/sample/5.png'  # 替换为你的图片路径
# image_path = './model/sample/7.png'
image_path = './build/img/0.png'
image = Image.open(image_path)
image = image.resize((28, 28))  # 调整图片大小为28x28
image = transform(image).unsqueeze(0)  # 添加批次维度

# 推理验证
with torch.no_grad():  # 禁用梯度计算以加速推理过程
    outputs = model(image)
    _, predicted = torch.max(outputs.data, 1)

print(f'The predicted digit is: {predicted.item()}')

# 显示图片
plt.imshow(image.squeeze().numpy(), cmap='gray')  # 使用 numpy() 方法将张量转换为 numpy 数组
plt.title(f'Predicted Digit: {predicted.item()}')
plt.show()

# 给一个用mnist验证一张图片是数字及的代码
# 在这个代码中：
#
# 我们定义了数据转换，将图像转换为灰度图、张量并进行归一化。
# 加载预训练的模型（假设模型保存在
# mnist_model.pth
# 文件中）。
# 加载并预处理单张图片，将其调整为28x28的大小，并应用相同的数据转换。
# 将模型设置为评估模式，并禁用梯度计算以加快推理速度。
# 通过模型进行推理，并输出预测的数字。
# 最后，使用Matplotlib显示图片及其预测结果。
# 请确保你有一个预训练的模型文件
# mnist_model.pth，以及你要验证的图片路径。
